'''
Created on Mar 1, 2020
Pytorch Implementation of LightGCN in
Xiangnan He et al. LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation
@author: Jianbai Ye (gusye@mail.ustc.edu.cn)

Design training and test process
'''
import world
import numpy as np
import torch
import utils
import dataloader
from pprint import pprint
from utils import timer
from time import time
from tqdm import tqdm
import model
import multiprocessing
from sklearn.metrics import roc_auc_score


CORES = multiprocessing.cpu_count() // 2


def BPR_train_original(dataset, recommend_model, loss_class, epoch, neg_k=1, w=None):
    Recmodel = recommend_model
    Recmodel.train()
    """
    bpr: utils.BPRLoss: 这是一个类型注解，指示变量bpr应该是一个BPRLoss类的实例
    = loss_class: 这是一个赋值语句，它将右侧的loss_class变量（预期是BPRLoss类的一个实例）赋值给左侧的bpr变量。
    """
    bpr: utils.BPRLoss = loss_class
    
    with timer(name="Sample"):
        #utils.UniformSample_original：[users,pos,neg] 每个用户随机选交互过的一个商品和没有交互过的一个商品
        S = utils.UniformSample_original(dataset)
    #users所有用户
    users = torch.Tensor(S[:, 0]).long()
    #每个用户对应随机一个交互商品
    posItems = torch.Tensor(S[:, 1]).long()
    # 每个用户对应随机一个没有交互商品
    negItems = torch.Tensor(S[:, 2]).long()

    users = users.to(world.device)
    posItems = posItems.to(world.device)
    negItems = negItems.to(world.device)
    #shuffle打乱
    users, posItems, negItems = utils.shuffle(users, posItems, negItems)
    total_batch = len(users) // world.config['bpr_batch_size'] + 1
    aver_loss = 0.
    for (batch_i,
         (batch_users,
          batch_pos,
          batch_neg)) in enumerate(utils.minibatch(users,
                                                   posItems,
                                                   negItems,
                                                   batch_size=world.config['bpr_batch_size'])):
        #stageOne即开始训练反向传播并返回损失值用于记录
        cri = bpr.stageOne(batch_users, batch_pos, batch_neg)
        aver_loss += cri
        if world.tensorboard:
            w.add_scalar(f'BPRLoss/BPR', cri, epoch * int(len(users) / world.config['bpr_batch_size']) + batch_i)
    #计算平均损失
    aver_loss = aver_loss / total_batch
    #记录时间信息
    time_info = timer.dict()
    timer.zero()
    #返回平均损失和时间信息
    return f"loss{aver_loss:.3f}-{time_info}"
    
    
def test_one_batch(X):
    #每个X包含两个元素  第一个是用户的预测评分最高的20个商品集合[100][20] 第二个是测试集用户交互过的商品列表[100][测试集实际交互的商品数量]
    #sorted_items用户预测评分最高的20个商品集合  [100][20]
    sorted_items = X[0].numpy()
    #groundTrue是测试集用户交互过的商品列表 其形状为[100][测试集实际交互的商品数量]
    groundTrue = X[1]
    #getLabel#获取预测数据和测试数据矩阵 返回[batch][k] batch为一个批次的大小 k为排前k个商品 在这里默认为[100][20] 每个元素命中为1.没命中为0.
    r = utils.getLabel(groundTrue, sorted_items)
    pre, recall, ndcg = [], [], []
    #k遍历worl.topks中的每个元素 默认里面只有一个元素20
    for k in world.topks:
        ret = utils.RecallPrecision_ATk(groundTrue, r, k)
        pre.append(ret['precision'])
        recall.append(ret['recall'])
        ndcg.append(utils.NDCGatK_r(groundTrue,r,k))
    return {'recall':np.array(recall), 
            'precision':np.array(pre), 
            'ndcg':np.array(ndcg)}
        
            
def Test(dataset, Recmodel, epoch, w=None, multicore=0):
    u_batch_size = world.config['test_u_batch_size']
    #设置参数类型
    dataset: utils.BasicDataset
    testDict: dict = dataset.testDict
    Recmodel: model.LightGCN
    # eval mode with no dropout
    #eval让模型进入评估模式 会自动关闭dropout和梯度更新
    Recmodel = Recmodel.eval()
    #这里的topk默认为20
    max_K = max(world.topks)
    if multicore == 1:
        pool = multiprocessing.Pool(CORES)
    #创建全0矩阵
    results = {'precision': np.zeros(len(world.topks)),
               'recall': np.zeros(len(world.topks)),
               'ndcg': np.zeros(len(world.topks))}
    with torch.no_grad():
        #testDict即测试数据集 其格式为{user: [items]} 其中user为用户 [items]为用户交互过的所有商品序列
        #testDict.keys()获取所有用户 [users_len]
        users = list(testDict.keys())
        try:
            assert u_batch_size <= len(users) / 10
        except AssertionError:
            #batch太大调小一点
            print(f"test_u_batch_size is too big for this dataset, try a small one {len(users) // 10}")
        users_list = []
        rating_list = []
        groundTrue_list = []
        # auc_record = []
        # ratings = []
        #计算batch的个数
        total_batch = len(users) // u_batch_size + 1
        #users:[users] u_batch_size即test_batch为100
        for batch_users in utils.minibatch(users, batch_size=u_batch_size):
            #allpos为训练集用户数据 dataset.getUserPosItems每次获取batch_users个用户的所有交互商品记录 allpos[batch_users][所有商品交互信息]
            allPos = dataset.getUserPosItems(batch_users)
            #testDict为测试集用户数据 获取batch_users个用户的所有商品交互记录
            #allPos groundTrue都为用户和商品的交互记录但前者是训练集数据 后者是测试集数据
            groundTrue = [testDict[u] for u in batch_users]
            #把batch_users转为tensor对象并交给gpu处理
            batch_users_gpu = torch.Tensor(batch_users).long()
            batch_users_gpu = batch_users_gpu.to(world.device)
            #获取用户预测评分 维度:[batch_users_gpu，item_len]
            rating = Recmodel.getUsersRating(batch_users_gpu)
            #rating = rating.cpu()
            exclude_index = []
            exclude_items = []
            for range_i, items in enumerate(allPos):
                #把训练集用户交互数据allPos改为如下显示
                #exclude_items=[7,8,9,7,8,9,2,3,3....]物品数据
                #exclude_index=[0,0,0,1,1,2,2,2,2....]用户数据
                exclude_index.extend([range_i] * len(items))
                exclude_items.extend(items)
            #这通常用于在计算排名时排除某些元素。排除掉用户训练集中用户访问过的商品编号(用户训练集中访问过的商品编号不回出现在测试集中)
            rating[exclude_index, exclude_items] = -(1<<10) #1<<10将数字1左移10位，产生一个二进制数，然后取其负值 -(1<<10)=-1024
            #torch.topk(): 这是PyTorch中的一个函数，用于从张量中提取最大的k个值及其索引。这里的max_K默认为20
            #第一个是最大的k个值，第二个是这些最大值的索引 我们只要下标 即rating_K为每个用户打分最高的k个商品序号集合
            _, rating_K = torch.topk(rating, k=max_K)
            rating = rating.cpu().numpy()
            # aucs = [ 
            #         utils.AUC(rating[i],
            #                   dataset, 
            #                   test_data) for i, test_data in enumerate(groundTrue)
            #     ]
            # auc_record.extend(aucs)
            #用于删除对象 我们只需要每个用户预测的最大k个评分即可 即rating_K
            del rating
            #将一个batch的用户加到users_list中
            users_list.append(batch_users)
            #rating_K.cpu()把rating_K转移到cpu
            rating_list.append(rating_K.cpu())
            #groundTrue为测试集用户交互过的商品
            groundTrue_list.append(groundTrue)
        assert total_batch == len(users_list)
        #评分表 和 测试集用户交互秒
        #rating_list[299][100][20] 一共299组 每组100个用户 每个用户取预测最高的20个商品序号列表
        # groundTrue_list [299][100][测试集实际交互的商品数量] 一共299组 每组100个用户 获取每个用户实际交互数量
        X = zip(rating_list, groundTrue_list)
        if multicore == 1:
            #则使用 pool.map() 方法并行地对列表 X 中的每个元素应用 test_one_batch 函数。
            pre_results = pool.map(test_one_batch, X)
        else:
            pre_results = []
            for x in X:
                #对列表 X 中的每个元素 x 单独调用 test_one_batch 函数
                #x一共运行batch次即299次 每个小x包含两个元素
                #x[0]是用户取预测最高的20个商品序号列表  [100][20]
                #x[1]是测试集用户交互过的商品列表       [100][测试集实际交互的商品数量]
                pre_results.append(test_one_batch(x))
        scale = float(u_batch_size/len(users))
        #计算recall precision ndcg
        for result in pre_results:
            results['recall'] += result['recall']
            results['precision'] += result['precision']
            results['ndcg'] += result['ndcg']
        results['recall'] /= float(len(users))
        results['precision'] /= float(len(users))
        results['ndcg'] /= float(len(users))
        # results['auc'] = np.mean(auc_record)
        #数据记录
        if world.tensorboard:
            w.add_scalars(f'Test/Recall@{world.topks}',
                          {str(world.topks[i]): results['recall'][i] for i in range(len(world.topks))}, epoch)
            w.add_scalars(f'Test/Precision@{world.topks}',
                          {str(world.topks[i]): results['precision'][i] for i in range(len(world.topks))}, epoch)
            w.add_scalars(f'Test/NDCG@{world.topks}',
                          {str(world.topks[i]): results['ndcg'][i] for i in range(len(world.topks))}, epoch)
        if multicore == 1:
            pool.close()
        print(results)
        return results
